from langchain.prompts import PromptTemplate,SystemMessagePromptTemplate,HumanMessagePromptTemplate,ChatPromptTemplate
from typing import List
from datetime import datetime
from langchain_core.messages.base import BaseMessage
from langchain_core.messages.chat import ChatMessageChunk
from langchain_core.messages.human import HumanMessage
from langchain_core.messages.ai import AIMessage

# system_prompt="楼盘资料\n{context}"
system_prompt="""当前时间是{time}。
文档内容：{context}

Use the following format:
thought: 思考过程，关联的文档内容
Final Answer: 依照文档内容作出回答

回答内容简洁,以搜索到的数据为准，逻辑清晰。不进行过多的描述。
"""

chat_prompt="{input}"

class BasePrompt:
    def __init__(self,sep='\n',system_prompt=system_prompt,chat_prompt=chat_prompt):
        self.system_prompt=system_prompt
        self.chat_prompt=chat_prompt
        self.sep=sep

    def get_history_func(self,histroy:list=[]):
        res = []
        for message in histroy:
            temp = ChatMessageChunk(content=message['content'],role="user" if message['user_type'] == 0 else "assistant")
            res.append(temp)
        return res
    
    def init_prompt(self,ai_prompt:str = '' , histroy:List[BaseMessage]=[]):
        systemMessage=SystemMessagePromptTemplate.from_template(template=ai_prompt + self.system_prompt)

        userMessage=HumanMessagePromptTemplate.from_template(template=self.chat_prompt,)
        histroy=self.get_history_func(histroy)
        history=[systemMessage]+histroy+[userMessage]
        prompt=ChatPromptTemplate.from_messages(history)
        self.prompt=prompt

        prompt=prompt.partial(time=self.get_current_time())
        return prompt
    
    def get_current_time(self)->str:
        """
        获取当前日期和时间
        """

        return datetime.now().strftime("%Y年%m月%d日%H:%M")